#include <stdio.h>
#include <stdint.h>
#include <fstream>
#include <vector>
#include <chrono>
#include <cmath>
#include <limits>
#include <tuple>
#include <omp.h>

#define DIV_UP(size, align_size) (((size) + (align_size) - 1) / (align_size))

static std::uint64_t now() {
  return std::chrono::duration_cast<std::chrono::milliseconds>
    (std::chrono::high_resolution_clock::now().time_since_epoch()).count();
}

void dumpbin(void *buf, size_t sz, const std::string &path) {
  std::ofstream ofs(path, std::ios::out | std::ios::trunc | std::ios::binary);
  if (ofs.is_open()){
    ofs.write((const char *)buf, sz);
    // printf("# dump %s\n", path.c_str());
  }else{
    printf("open %s failed!", path.c_str());
  }
}

template<typename T>
void memset(void *buf, const T val, int sz) {
  for(int id = 0; id < sz; ++id){
    ((T*)buf)[id] = (id%2) ? (val+id*0.01f) : (val-id*0.01f);
  }
}

//safe softmax
template<typename T>
void softmax(int m, T *src, T *dst) {
  //normalization term
  float norm_term = 0.f;
  for(int i = 0; i < m; ++i){
    norm_term += expf(src[i]);
  }
  //safe softmax
  printf("%f\n", norm_term);
  for(int i = 0; i < m; ++i){
    dst[i] = expf(src[i])/norm_term;
  }
}

//safe softmax
template<typename T>
void softmax_safe(int m, T *src, T *dst) {
  //max_term x
  float max_term = std::numeric_limits<float>::min();
  for(int i = 0; i < m; ++i){
    max_term = std::max(src[i], max_term);
  }
  //normalization term
  float norm_term = 0.f;
  for(int i = 0; i < m; ++i){
    norm_term += expf(src[i] - max_term);
  }
  //safe softmax
  for(int i = 0; i < m; ++i){
    dst[i] = expf(src[i] - max_term)/norm_term;
  }
}

/**
 * @brief 并行，在线，安全，softmax
 * @see 参考文章：
 * docs/1805.02867v2.pdf
 */

//online, safe softmax
template<typename T>
void softmax_safe_online(int m, T *src, T *dst) {
  //max_term, norm_term
  float max_term = std::numeric_limits<float>::min(), norm_term = 0;
  for(int i = 0; i < m; ++i){
    float pre_max_term = max_term;
    max_term = std::max(src[i], max_term);
    norm_term = norm_term*expf(pre_max_term-max_term)+expf(src[i]-max_term);
  }
  //safe softmax
  for(int i = 0; i < m; ++i){
    dst[i] = expf(src[i] - max_term)/norm_term;
  }
}

//parallel, online, safe softmax
void softmax_safe_online_parallel(int m, float *src, float *dst) {
  
  auto reduce_op = [](const std::tuple<float, float> &left, const std::tuple<float, float> &right)-> std::tuple<float, float>{
    float m_i, d_i, m_j, d_j, m_ret, d_ret;
    std::tie(m_i, d_i) = left;
    std::tie(m_j, d_j) = right;
    m_ret = std::max(m_i, m_j);
    d_ret = d_i*expf(m_i-m_ret)+d_j*expf(m_j-m_ret);
    return std::make_tuple(m_ret, d_ret);
  };
  const int nb = omp_get_num_procs();
  int szBlock = DIV_UP(m, nb);
  std::vector<std::tuple<float, float>> terms;
  terms.resize(nb);

  omp_set_num_threads(nb); // 设置线程数

  //parallel compute block term
  #pragma omp parallel for
  for(int iBLock = 0; iBLock < nb; ++iBLock){
    int iStart = iBLock*szBlock;
    float max_term = std::numeric_limits<float>::min(), norm_term = 0.f;
    for(int i = iStart; (i < iStart+szBlock) && (i < m); ++i){
      std::tie(max_term, norm_term) = reduce_op({max_term, norm_term}, {src[i], 1});
    }
    terms[iBLock] = {max_term, norm_term};
  }
  //reduce all terms
  float max_term = std::numeric_limits<float>::min(), norm_term = 0.f;
  for(int iTerm = 0; iTerm < nb; ++iTerm){
    std::tie(max_term, norm_term) = reduce_op({max_term, norm_term}, terms[iTerm]);
  }

  //parallel compute block softmax
  #pragma omp parallel for
  for(int iBLock = 0; iBLock < nb; ++iBLock){
    int iStart = iBLock*szBlock;
    for(int i = iStart; (i < iStart+szBlock) && (i < m); ++i){
      dst[i] = expf(src[i] - max_term)/norm_term;
    }
  }
}

struct Case{
  int m;
  float src_mask;
  std::string src_path;
  std::string dst_path;
  int loop;
  float avg_latency;
};

int main(int argc, char *argv[]) {
  std::vector<Case> cases;

#ifdef EVAL
  cases.push_back({4, 1.0f, "data/c0_src.bin",  "data/c0_dst.bin", 1});
  cases.push_back({8, 1.0f, "data/c1_src.bin",  "data/c1_dst.bin", 1});
  cases.push_back({1000, 1.0f, "data/c2_src.bin",  "data/c2_dst.bin", 1});
#endif

#ifdef PERF
  float m_start = 100000;
  int loop = 1000;

  for(int m_id = m_start; m_id <= 10*m_start; m_id+=m_start){
      cases.push_back({m_id, 1.0f, "data/tmp_src.bin",  "data/tmp_dst.bin", loop});
  }
#endif

  for(int id = 0; id < cases.size(); ++id){
    #ifdef PERF
    if(id > 0){
        cases[id].loop = std::max(1, int(1000.0f/cases[id-1].avg_latency));
    }
    #endif
    auto c = cases[id];
    const int m = c.m;
    std::string src_path = c.src_path, dst_path = c.dst_path;
    float src_mask = c.src_mask;

    float *src = new float[m];
    float *dst = new float[m];
    memset(src, src_mask, m);
    // memset(dst, 0, m);
    
    int start = now();
    for(int i = 0; i < c.loop; ++i){
      // softmax(m, src, dst);
      // softmax_safe(m, src, dst);
      // softmax_safe_online(m, src, dst);
      softmax_safe_online_parallel(m, src, dst);
    }
    int end = now();
    cases[id].avg_latency =float(end-start)/float(c.loop);

    #ifdef EVAL
    dumpbin(src, m*sizeof(float), src_path);
    dumpbin(dst, m*sizeof(float), dst_path);
    // printf("result %0.2f \n", dst[0]);
    printf("{\"m\": %d, \"src_path\": \"%s\", \"dst_path\": \"%s\"}\n", m, src_path.c_str(), dst_path.c_str());
    #endif

    #ifdef PERF
    {
      Case c = cases[id];
      const int m = c.m, loop = c.loop;
      float avg_latency = c.avg_latency;
      printf("  m %d loop %d avg_latency %0.4f ms  \n", m, loop, avg_latency);
    }
    #endif

    delete []src;
    delete []dst;
  }
  return 0;
}